#include<iostream>
#include<queue>
using namespace std;

/*
extern "C" int func_name(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);
nparam (int): 输入输出总数。比如算子有2个输入，1个输出，则nparam的值为3。
params (void **): 输入输出指针数组。比如算子有2个输入，1个输出，params[0]指向第一个输入数据，params[1]指向第二个输入数据，params[2]指向输出数据
ndims (int *): 输入输出shape维度数组。比如params[i]是个shape[1024, 1024]的张量，则ndims[i]的值为2。
shapes (int64_t **): 输入输出shape数组。比如params[i]是个shape[1024, 1024]的张量，则shapes[i][0]的值为1024，shapes[i][1]的值为1024。
dtypes (const char **): 输入输出数据类型数组。dtypes里的元素取值可为：
”float32”, “float16”, “float”, “float64”, “int”, “int8”, “int16”, “int32”, “int64”, “uint”, “uint8”, “uint16”, “uint32”, “uint64”, “bool”。
stream (void *): CUDA流指针，仅定义GPU算子实现时需要。
extra (void *): 用于后续扩展。
*/
/*
shapes=([1,1,712,1072],[1,1,712,1072],[1,1,712,1072])
op_test=ops.Custom("./cal_RE_PR.so:CustomTest",out_shape=[1,1,712,1072],out_dtype=ms.float32,func_type="aot",reg_info=custom_op_test)
*/
extern "C" int CustomTest(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra){
    //nparam=4
    //nidms:[4,4,4,1]
    //shapes:shape[0][0-3]=[1,1,712,1072]
    //shapes:shape[1][0-3]=[1,1,712,1072]
    //shapes:shape[2][0-3]=[1,1,712,1072]
    //shapes:shape[3][0]=[2]
    if (nparam!=4) return 1;
    float *data=static_cast<float *>(params[0]);
    float *label=static_cast<float *>(params[1]);
    int *isVisited=static_cast<int *>(params[2]);
    float *output=static_cast<float *>(params[3]);
    
    int length=shapes[0][2]*shapes[0][3];//712*1072
    int w=shapes[0][2];
    int h=shapes[0][3];
    float r=0.5;//判别系数,表示连通域判定的严格程度
    int TP=0;
    int FP=0;
    int FN=0;
    queue<int>queue;//BFS算法所需的队列

    /*
    data是模型训练后得到的张量/图像,未做处理,包含[0,255]的像素点
    label是标签图像,其中病变点像素值为255,背景点为0
    */
    for(int i=0;i<length;i++){
        if(data[i]!=255||isVisited[i]==1)continue;
        else{
            //data[i]==255&&isVisited[i]==0
            //此时data[i]=255且尚未访问过,进行BFS搜索
            int numTotal=0;//像素总数
            int numOverlap=0;//重叠区域像素数
            queue.push(i);//当前坐标入队
            isVisited[i]=1;//标记已访问过
            if(label[i]==255)numOverlap++;//判断data[i]是否为TP
            while(!queue.empty()){
                int temp=queue.front();//较早入队的坐标退出,并获取其值
                //当前坐标的上下左右节点入队(如果有且符合条件,即不越过边界)
                //一维数组的原二维数组坐标：array[i]->array[i/w][i-(int)(i/w)*w]
                //当前坐标的上下左右坐标为temp-w、temp+w、temp-1、temp+1
                int x=temp/w;
                int y=temp-w*(int)(temp/w);
                //上边界
                if(temp>=w&&data[temp-w]==255&&isVisited[temp]==0){
                    queue.push(temp-w);
                    isVisited[temp-w]=1;
                    numTotal++;
                    if(label[temp-w]==0)numOverlap++;
                }
                //下边界
                if(temp<=length-1-w&&data[temp+w]==255&&isVisited[temp]==0){
                    queue.push(temp+w);
                    isVisited[temp+w]=1;
                    numTotal++;
                    if(label[temp+w]==0)numOverlap++;
                }
                //左边界
                if(temp%w!=0&&data[temp-1]==255&&isVisited[temp]==0){
                    queue.push(temp-1);
                    isVisited[temp-1]=1;
                    numTotal++;
                    if(label[temp-1]==0)numOverlap++;
                }
                //右边界
                if((temp+1)%w!=0&&data[temp+1]==255&&isVisited[temp]==0){
                    queue.push(temp+1);
                    isVisited[temp+1]=1;
                    numTotal++;
                    if(label[temp+1]==0)numOverlap++;
                }
                }
                if((numOverlap/numTotal)>r)TP+=numTotal;//全部区域作为标记区
                else{//仅重叠区域作为标记区,其余区域作为假阳性区
                    TP+=numOverlap;
                    FP+=numTotal-numOverlap;
                }   
        }
    }

    for(int i=0;i<length;i++)isVisited[i]=0;//判别矩阵恢复默认
    while(!queue.empty())queue.pop();//清空queue

    for(int i=0;i<length;i++){
        if(label[i]!=255||isVisited[i]==1)continue;
        else{
            //label[i]==255&&isVisited[i]==0
            //此时label[i]=255且尚未访问过,进行BFS搜索
            int numTotal=0;//像素总数
            int numOverlap=0;//重叠区域像素数
            queue.push(i);//当前坐标入队
            isVisited[i]=1;//标记已访问过
            if(data[i]==255)numOverlap++;//判断label[i]是否为TP
            while(!queue.empty()){
                int temp=queue.front();//较早入队的坐标退出,并获取其值
                //当前坐标的上下左右节点入队(如果有且符合条件)
                //一维数组的原二维数组坐标：array[i]->array[i/w][i-(int)(i/w)*w]
                //当前坐标的上下左右坐标为temp-w、temp+w、temp-1、temp+1
                int x=temp/w;
                int y=temp-w*(int)(temp/w);
                //上边界
                if(temp>=w&&label[temp-w]==255&&isVisited[temp]==0){
                    queue.push(temp-w);
                    isVisited[temp-w]=1;
                    numTotal++;
                    if(data[temp-w]==0)numOverlap++;
                }
                //下边界
                if(temp<=length-1-w&&label[temp+w]==255&&isVisited[temp]==0){
                    queue.push(temp+w);
                    isVisited[temp+w]=1;
                    numTotal++;
                    if(data[temp+w]==0)numOverlap++;
                }
                //左边界
                if(temp%w!=0&&label[temp-1]==255&&isVisited[temp]==0){
                    queue.push(temp-1);
                    isVisited[temp-1]=1;
                    numTotal++;
                    if(data[temp-1]==0)numOverlap++;
                }
                //右边界
                if((temp+1)%w!=0&&label[temp+1]==255&&isVisited[temp]==0){
                    queue.push(temp+1);
                    isVisited[temp+1]=1;
                    numTotal++;
                    if(data[temp+1]==0)numOverlap++;
                }
            }
            if((numOverlap/numTotal)>r)TP+=numTotal;//全部区域作为标记区
            else FN+=numTotal=numOverlap;//仅重叠区域作为标记区,其余区域作为假阳性区
        }
    }
    float RE=TP/(TP+FN);
    float PR=TP/(TP/FP);
    output[0]=RE;
    output[1]=PR;
    return 0;//仅作为判断函数是否执行正确
}